package db

import (
	"database/sql"
	"encoding/json"
	"log"
	"strings"
	"time"

	"github.com/xfio/internal/models"
	_ "github.com/mattn/go-sqlite3"
)

type DB struct {
	*sql.DB
}

func NewDB(filepath string) (*DB, error) {
	db, err := sql.Open("sqlite3", filepath + "?_busy_timeout=10000&_journal_mode=WAL")
	if err != nil {
		return nil, err
	}

	// 创建workers表
	statement1 := `
	CREATE TABLE IF NOT EXISTS workers (
		id TEXT PRIMARY KEY,
		name TEXT NOT NULL,
		ip_address TEXT NOT NULL,
		status TEXT DEFAULT 'offline',
		heartbeat_status TEXT DEFAULT 'healthy',
		last_heartbeat DATETIME,
		consecutive_heartbeat_failures INTEGER DEFAULT 0,
		consecutive_heartbeat_success INTEGER DEFAULT 0,
		created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
		system_info TEXT,
		disk_info TEXT,
		fio_version TEXT DEFAULT '',
		fio_status TEXT DEFAULT 'stopped'
	);
	`

	_, err = db.Exec(statement1)
	if err != nil {
		return nil, err
	}
	
	// 为现有表添加 fio_version 字段（如果不存在）
	_, err = db.Exec("ALTER TABLE workers ADD COLUMN fio_version TEXT DEFAULT ''")
	if err != nil {
		// 忽略错误，因为字段可能已经存在
		// log.Printf("Note: Column fio_version may already exist: %v", err)
	}

	// 为现有表添加 fio_status 字段（如果不存在）
	_, err = db.Exec("ALTER TABLE workers ADD COLUMN fio_status TEXT DEFAULT 'stopped'")
	if err != nil {
		// 忽略错误，因为字段可能已经存在
		// log.Printf("Note: Column fio_status may already exist: %v", err)
	}

	// 创建test_cases表
	statement2 := `
	CREATE TABLE IF NOT EXISTS test_cases (
		id TEXT PRIMARY KEY,
		name TEXT NOT NULL,
		description TEXT,
		fio_config TEXT NOT NULL,
		created_at DATETIME DEFAULT CURRENT_TIMESTAMP
	);
	`

	_, err = db.Exec(statement2)
	if err != nil {
		return nil, err
	}

	// 创建tasks表
	statement3 := `
	CREATE TABLE IF NOT EXISTS tasks (
		id TEXT PRIMARY KEY,
		test_case_id TEXT,
		worker_id TEXT,
		worker_ids TEXT,
		status TEXT DEFAULT 'pending',
		start_time DATETIME,
		end_time DATETIME,
		result TEXT,
		task_config_override TEXT,
		created_at DATETIME DEFAULT CURRENT_TIMESTAMP
	);
	`

	_, err = db.Exec(statement3)
	if err != nil {
		return nil, err
	}

	// 尝试为旧数据库添加 task_config_override 字段与 worker_ids 字段（如果不存在）
	_, err = db.Exec("ALTER TABLE tasks ADD COLUMN task_config_override TEXT")
	if err != nil {
		log.Printf("Note: Column task_config_override may already exist: %v", err)
	}
	_, err = db.Exec("ALTER TABLE tasks ADD COLUMN worker_ids TEXT")
	if err != nil {
		log.Printf("Note: Column worker_ids may already exist: %v", err)
	}
	_, err = db.Exec("ALTER TABLE tasks ADD COLUMN name TEXT")
	if err != nil {
		log.Printf("Note: Column name may already exist: %v", err)
	}

	return &DB{db}, nil
}

func (db *DB) RegisterWorker(worker models.Worker) error {
	// 将系统信息和磁盘信息转换为JSON字符串
	var systemInfoJSON, diskInfoJSON string
	var err error
	
	if worker.SystemInfo != nil {
		systemInfoBytes, err := json.Marshal(worker.SystemInfo)
		if err != nil {
			return err
		}
		systemInfoJSON = string(systemInfoBytes)
	}
	
	if worker.DiskInfo != nil {
		diskInfoBytes, err := json.Marshal(worker.DiskInfo)
		if err != nil {
			return err
		}
		diskInfoJSON = string(diskInfoBytes)
	}

	// 首先检查是否已存在相同名称和IP的Worker
	var existingID string
	
	err = db.QueryRow("SELECT id FROM workers WHERE name = ? AND ip_address = ?", worker.Name, worker.IPAddress).Scan(&existingID)
	
	// 如果存在相同名称和IP的Worker，则更新现有记录而不是创建新记录
	if err == nil {
		// Worker已存在，更新记录
		statement := `
		UPDATE workers 
		SET status = ?, heartbeat_status = ?, last_heartbeat = ?, consecutive_heartbeat_success = ?, system_info = ?, disk_info = ?, fio_version = ?, fio_status = ?
		WHERE id = ?
		`
		
		_, err := db.Exec(statement, worker.Status, worker.HeartbeatStatus, worker.LastHeartbeat, worker.ConsecutiveHeartbeatSuccess, systemInfoJSON, diskInfoJSON, worker.FIOVersion, worker.FIOStatus, existingID)
		return err
	} else if err != sql.ErrNoRows {
		// 其他数据库错误
		return err
	}
	
	// Worker不存在，创建新记录
	statement := `
	INSERT INTO workers (id, name, ip_address, status, heartbeat_status, last_heartbeat, consecutive_heartbeat_success, system_info, disk_info, fio_version, fio_status)
	VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	`

	_, err = db.Exec(statement, worker.ID, worker.Name, worker.IPAddress, worker.Status, worker.HeartbeatStatus, worker.LastHeartbeat, worker.ConsecutiveHeartbeatSuccess, systemInfoJSON, diskInfoJSON, worker.FIOVersion, worker.FIOStatus)
	return err
}

func (db *DB) GetWorkers() ([]models.Worker, error) {
	rows, err := db.Query("SELECT id, name, ip_address, status, heartbeat_status, last_heartbeat, created_at, consecutive_heartbeat_failures, consecutive_heartbeat_success, system_info, disk_info, fio_version, fio_status FROM workers")
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var workers []models.Worker
	for rows.Next() {
		var w models.Worker
		var systemInfoStr, diskInfoStr, fioVersionStr, fioStatusStr sql.NullString
		var lastHeartbeat, createdAt sql.NullTime
		err := rows.Scan(&w.ID, &w.Name, &w.IPAddress, &w.Status, &w.HeartbeatStatus, &lastHeartbeat, &createdAt, &w.ConsecutiveHeartbeatFailures, &w.ConsecutiveHeartbeatSuccess, &systemInfoStr, &diskInfoStr, &fioVersionStr, &fioStatusStr)
		if err != nil {
			return nil, err
		}
		
		// 处理时间字段
		if lastHeartbeat.Valid {
			w.LastHeartbeat = lastHeartbeat.Time
		}
		if createdAt.Valid {
			w.CreatedAt = createdAt.Time
		}
		
		// 解析系统信息
		if systemInfoStr.Valid && systemInfoStr.String != "" {
			var systemInfo models.SystemInfo
			if err := json.Unmarshal([]byte(systemInfoStr.String), &systemInfo); err == nil {
				w.SystemInfo = &systemInfo
			}
		}
		
		// 解析磁盘信息
		if diskInfoStr.Valid && diskInfoStr.String != "" {
			var diskInfo []models.DiskInfo
			if err := json.Unmarshal([]byte(diskInfoStr.String), &diskInfo); err == nil {
				w.DiskInfo = diskInfo
			}
		}
		
		// 设置 fio 版本信息
		if fioVersionStr.Valid {
			w.FIOVersion = fioVersionStr.String
		}
		
		// 设置 fio 状态信息
		if fioStatusStr.Valid {
			w.FIOStatus = fioStatusStr.String
		}
		
		workers = append(workers, w)
	}

	return workers, nil
}

func (db *DB) GetWorker(id string) (*models.Worker, error) {
	var w models.Worker
	var systemInfoStr, diskInfoStr, fioVersionStr, fioStatusStr sql.NullString
	var lastHeartbeat, createdAt sql.NullTime
	err := db.QueryRow("SELECT id, name, ip_address, status, heartbeat_status, last_heartbeat, created_at, consecutive_heartbeat_failures, consecutive_heartbeat_success, system_info, disk_info, fio_version, fio_status FROM workers WHERE id = ?", id).
		Scan(&w.ID, &w.Name, &w.IPAddress, &w.Status, &w.HeartbeatStatus, &lastHeartbeat, &createdAt, &w.ConsecutiveHeartbeatFailures, &w.ConsecutiveHeartbeatSuccess, &systemInfoStr, &diskInfoStr, &fioVersionStr, &fioStatusStr)
	if err != nil {
		return nil, err
	}
	
	// 处理时间字段
	if lastHeartbeat.Valid {
		w.LastHeartbeat = lastHeartbeat.Time
	}
	if createdAt.Valid {
		w.CreatedAt = createdAt.Time
	}
	
	// 解析系统信息
	if systemInfoStr.Valid && systemInfoStr.String != "" {
		var systemInfo models.SystemInfo
		if err := json.Unmarshal([]byte(systemInfoStr.String), &systemInfo); err == nil {
			w.SystemInfo = &systemInfo
		}
	}
	
	// 解析磁盘信息
	if diskInfoStr.Valid && diskInfoStr.String != "" {
		var diskInfo []models.DiskInfo
		if err := json.Unmarshal([]byte(diskInfoStr.String), &diskInfo); err == nil {
			w.DiskInfo = diskInfo
		}
	}
	
	// 设置 fio 版本信息
	if fioVersionStr.Valid {
		w.FIOVersion = fioVersionStr.String
	}
	
	// 设置 fio 状态信息
	if fioStatusStr.Valid {
		w.FIOStatus = fioStatusStr.String
	}

	return &w, nil
}

func (db *DB) UpdateWorkerHeartbeat(id string) error {
	// 获取当前心跳失败次数
	var currentFailures int
	err := db.QueryRow("SELECT consecutive_heartbeat_failures FROM workers WHERE id = ?", id).
		Scan(&currentFailures)
	if err != nil {
		return err
	}

	statement := `
	UPDATE workers 
	SET status = 'online', last_heartbeat = ?, consecutive_heartbeat_success = consecutive_heartbeat_success + 1, consecutive_heartbeat_failures = 0
	WHERE id = ?
	`

	_, err = db.Exec(statement, time.Now(), id)
	if err != nil {
		return err
	}

	// 检查是否连续成功达到3次
	var currentSuccess int
	err = db.QueryRow("SELECT consecutive_heartbeat_success FROM workers WHERE id = ?", id).
		Scan(&currentSuccess)
	if err != nil {
		return err
	}

	if currentSuccess >= 3 {
		// 更新状态为健康，重置成功计数
		updateStatement := `
		UPDATE workers
		SET heartbeat_status = 'healthy', consecutive_heartbeat_success = 0
		WHERE id = ?
		`
		_, err = db.Exec(updateStatement, id)
		if err != nil {
			return err
		}
	}

	return nil
}

func (db *DB) DeleteWorker(id string) error {
	statement := `DELETE FROM workers WHERE id = ?`
	_, err := db.Exec(statement, id)
	return err
}

func (db *DB) CreateTestCase(testCase models.TestCase) error {
	fioConfigJSON, err := json.Marshal(testCase.FIOConfig)
	if err != nil {
		return err
	}

	statement := `
	INSERT INTO test_cases (id, name, description, fio_config)
	VALUES (?, ?, ?, ?)
	`

	_, err = db.Exec(statement, testCase.ID, testCase.Name, testCase.Description, string(fioConfigJSON))
	return err
}

func (db *DB) GetTestCases() ([]models.TestCase, error) {
	rows, err := db.Query("SELECT id, name, description, fio_config, created_at FROM test_cases")
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var testCases []models.TestCase
	for rows.Next() {
		var tc models.TestCase
		var fioConfigStr string
		err := rows.Scan(&tc.ID, &tc.Name, &tc.Description, &fioConfigStr, &tc.CreatedAt)
		if err != nil {
			return nil, err
		}

		err = json.Unmarshal([]byte(fioConfigStr), &tc.FIOConfig)
		if err != nil {
			return nil, err
		}

		testCases = append(testCases, tc)
	}

	return testCases, nil
}

// GetTestCasesByName performs a case-insensitive LIKE search on test case names.
// Passing empty query returns all test cases (same as GetTestCases).
func (db *DB) GetTestCasesByName(q string, limit int) ([]models.TestCase, error) {
	q = strings.TrimSpace(q)
	// reuse base query
	var rows *sql.Rows
	var err error
	if q == "" {
		if limit > 0 {
			rows, err = db.Query("SELECT id, name, description, fio_config, created_at FROM test_cases LIMIT ?", limit)
		} else {
			rows, err = db.Query("SELECT id, name, description, fio_config, created_at FROM test_cases")
		}
	} else {
		pattern := "%" + strings.ToLower(q) + "%"
		if limit > 0 {
			rows, err = db.Query("SELECT id, name, description, fio_config, created_at FROM test_cases WHERE lower(name) LIKE ? LIMIT ?", pattern, limit)
		} else {
			rows, err = db.Query("SELECT id, name, description, fio_config, created_at FROM test_cases WHERE lower(name) LIKE ?", pattern)
		}
	}
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var testCases []models.TestCase
	for rows.Next() {
		var tc models.TestCase
		var fioConfigStr string
		err := rows.Scan(&tc.ID, &tc.Name, &tc.Description, &fioConfigStr, &tc.CreatedAt)
		if err != nil {
			return nil, err
		}
		if err = json.Unmarshal([]byte(fioConfigStr), &tc.FIOConfig); err != nil {
			return nil, err
		}
		testCases = append(testCases, tc)
	}
	return testCases, nil
}

func (db *DB) GetTestCase(id string) (*models.TestCase, error) {
	var tc models.TestCase
	var fioConfigStr string
	err := db.QueryRow("SELECT id, name, description, fio_config, created_at FROM test_cases WHERE id = ?", id).
		Scan(&tc.ID, &tc.Name, &tc.Description, &fioConfigStr, &tc.CreatedAt)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal([]byte(fioConfigStr), &tc.FIOConfig)
	if err != nil {
		return nil, err
	}

	return &tc, nil
}

func (db *DB) UpdateTestCase(testCase models.TestCase) error {
	fioConfigJSON, err := json.Marshal(testCase.FIOConfig)
	if err != nil {
		return err
	}

	statement := `
	UPDATE test_cases
	SET name = ?, description = ?, fio_config = ?
	WHERE id = ?
	`

	_, err = db.Exec(statement, testCase.Name, testCase.Description, string(fioConfigJSON), testCase.ID)
	return err
}

func (db *DB) DeleteTestCase(id string) error {
	statement := `DELETE FROM test_cases WHERE id = ?`
	_, err := db.Exec(statement, id)
	return err
}

func (db *DB) BatchDeleteTestCases(ids []string) (int64, error) {
	if len(ids) == 0 {
		return 0, nil
	}
	
	// 构建SQL查询，使用占位符避免SQL注入
	placeholders := make([]string, len(ids))
	args := make([]interface{}, len(ids))
	for i, id := range ids {
		placeholders[i] = "?"
		args[i] = id
	}
	
	query := "DELETE FROM test_cases WHERE id IN (" + makePlaceholders(len(ids)) + ")"
	
	result, err := db.Exec(query, args...)
	if err != nil {
		return 0, err
	}
	
	deletedCount, err := result.RowsAffected()
	return deletedCount, err
}

// makePlaceholders 创建SQL查询中的占位符
func makePlaceholders(count int) string {
	if count <= 0 {
		return ""
	}
	
	placeholders := make([]string, count)
	for i := range placeholders {
		placeholders[i] = "?"
	}
	return strings.Join(placeholders, ",")
}

func (db *DB) AssignTask(task models.Task) error {
	var configOverrideJSON string
	var workerIDsJSON string
	if task.TaskConfigOverride != nil {
		configOverrideBytes, err := json.Marshal(task.TaskConfigOverride)
		if err != nil {
			return err
		}
		configOverrideJSON = string(configOverrideBytes)
	} else {
		configOverrideJSON = ""
	}

	if task.WorkerIDs != nil && len(task.WorkerIDs) > 0 {
		b, _ := json.Marshal(task.WorkerIDs)
		workerIDsJSON = string(b)
	} else {
		workerIDsJSON = ""
	}

	statement := `
	INSERT INTO tasks (id, name, test_case_id, worker_id, worker_ids, status, task_config_override)
	VALUES (?, ?, ?, ?, ?, ?, ?)
	`

	_, err := db.Exec(statement, task.ID, task.Name, task.TestCaseID, task.WorkerID, workerIDsJSON, task.Status, configOverrideJSON)
	return err
}

func (db *DB) GetTasks() ([]models.Task, error) {
	rows, err := db.Query(`
		SELECT id, name, test_case_id, worker_id, worker_ids, status, start_time, end_time, result, task_config_override, created_at
		FROM tasks
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var tasks []models.Task
	for rows.Next() {
		var t models.Task
		var resultStr sql.NullString
		var configOverrideStr sql.NullString
		var startTime sql.NullTime
		var endTime sql.NullTime
		var nameStr sql.NullString
		
			var workerIDsStr sql.NullString
			err := rows.Scan(&t.ID, &nameStr, &t.TestCaseID, &t.WorkerID, &workerIDsStr, &t.Status, &startTime, &endTime, &resultStr, &configOverrideStr, &t.CreatedAt)
		if err != nil {
			log.Printf("Error scanning task row: %v", err)
			return nil, err
		}

		if nameStr.Valid {
			t.Name = nameStr.String
		}

		if startTime.Valid {
			t.StartTime = &startTime.Time
		}
		if endTime.Valid {
			t.EndTime = &endTime.Time
		}
		if resultStr.Valid {
			err = json.Unmarshal([]byte(resultStr.String), &t.Result)
			if err != nil {
				return nil, err
			}
		}
		if configOverrideStr.Valid && configOverrideStr.String != "" {
			err = json.Unmarshal([]byte(configOverrideStr.String), &t.TaskConfigOverride)
			if err != nil {
				return nil, err
			}
		}

		if workerIDsStr.Valid && workerIDsStr.String != "" {
			var ids []string
			if err := json.Unmarshal([]byte(workerIDsStr.String), &ids); err == nil {
				t.WorkerIDs = ids
			}
		}

		tasks = append(tasks, t)
	}

	return tasks, nil
}

func (db *DB) GetTask(id string) (*models.Task, error) {
	var t models.Task
	var resultStr sql.NullString
	var configOverrideStr sql.NullString
	var workerIDsStr sql.NullString
	var startTime sql.NullTime
	var endTime sql.NullTime
	var nameStr sql.NullString
	
		err := db.QueryRow(`
			SELECT id, name, test_case_id, worker_id, worker_ids, status, start_time, end_time, result, task_config_override, created_at
			FROM tasks
			WHERE id = ?
		`, id).Scan(&t.ID, &nameStr, &t.TestCaseID, &t.WorkerID, &workerIDsStr, &t.Status, &startTime, &endTime, &resultStr, &configOverrideStr, &t.CreatedAt)
	if err != nil {
		return nil, err
	}

	if nameStr.Valid {
		t.Name = nameStr.String
	}

	if startTime.Valid {
		t.StartTime = &startTime.Time
	}
	if endTime.Valid {
		t.EndTime = &endTime.Time
	}
	if resultStr.Valid {
		err = json.Unmarshal([]byte(resultStr.String), &t.Result)
		if err != nil {
			return nil, err
		}
	}
	if configOverrideStr.Valid && configOverrideStr.String != "" {
		err = json.Unmarshal([]byte(configOverrideStr.String), &t.TaskConfigOverride)
		if err != nil {
			return nil, err
		}
	}

	return &t, nil
}

func (db *DB) UpdateTaskStatus(id, status string) error {
	statement := `
	UPDATE tasks
	SET status = ?, start_time = ?
	WHERE id = ?
	`

	_, err := db.Exec(statement, status, time.Now(), id)
	return err
}

// StopTask 将正在运行的任务标记为已停止/待处理（由控制器调用以响应停止请求）
func (db *DB) StopTask(id string) error {
	// 我们只将状态重置为 'pending'，并保留 start_time/end_time 以供审计（如需清除可调整）
	statement := `
	UPDATE tasks
	SET status = 'pending'
	WHERE id = ?
	`

	_, err := db.Exec(statement, id)
	return err
}

func (db *DB) CompleteTask(id string, result interface{}) error {
	resultJSON, err := json.Marshal(result)
	if err != nil {
		return err
	}

	statement := `
	UPDATE tasks
	SET status = 'completed', end_time = ?, result = ?
	WHERE id = ?
	`

	_, err = db.Exec(statement, time.Now(), string(resultJSON), id)
	return err
}

func (db *DB) GetPendingTasksForWorker(workerID string) ([]models.Task, error) {
	// 支持两种场景：
	// 1) 传统的 worker_id 字段匹配
	// 2) 新增的 worker_ids JSON 字段中包含该 worker
	likePattern := "%\"" + workerID + "\"%" // 匹配 JSON 字符串中的 "id"
	rows, err := db.Query(`
		SELECT t.id, t.test_case_id, t.worker_id, t.worker_ids, t.status, t.created_at, t.task_config_override
		FROM tasks t
		WHERE (t.worker_id = ? OR t.worker_ids LIKE ?) AND t.status = 'pending'
	`, workerID, likePattern)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var tasks []models.Task
	for rows.Next() {
		var t models.Task
		var configOverrideStr sql.NullString
		var workerIDsStr sql.NullString
		err := rows.Scan(&t.ID, &t.TestCaseID, &t.WorkerID, &workerIDsStr, &t.Status, &t.CreatedAt, &configOverrideStr)
		if err != nil {
			return nil, err
		}
		
		if configOverrideStr.Valid && configOverrideStr.String != "" {
			err = json.Unmarshal([]byte(configOverrideStr.String), &t.TaskConfigOverride)
			if err != nil {
				return nil, err
			}
		}

		if workerIDsStr.Valid && workerIDsStr.String != "" {
			var ids []string
			if err := json.Unmarshal([]byte(workerIDsStr.String), &ids); err == nil {
				t.WorkerIDs = ids
			}
		}
		
		tasks = append(tasks, t)
	}

	return tasks, nil
}

func (db *DB) GetCompletedTasksByTestCase(testCaseID string) ([]models.Task, error) {
	rows, err := db.Query(`
		SELECT id, worker_id, result
		FROM tasks
		WHERE test_case_id = ? AND status = 'completed'
	`, testCaseID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var tasks []models.Task
	for rows.Next() {
		var t models.Task
		var resultStr sql.NullString
		err := rows.Scan(&t.ID, &t.WorkerID, &resultStr)
		if err != nil {
			return nil, err
		}

		if resultStr.Valid {
			err = json.Unmarshal([]byte(resultStr.String), &t.Result)
			if err != nil {
				return nil, err
			}
		}

		tasks = append(tasks, t)
	}

	return tasks, nil
}

func (db *DB) DeleteTask(id string) error {
	_, err := db.Exec("DELETE FROM tasks WHERE id = ?", id)
	return err
}

// BatchDeleteTasks 批量删除任务，返回删除的数量
func (db *DB) BatchDeleteTasks(ids []string) (int64, error) {
	if len(ids) == 0 {
		return 0, nil
	}

	placeholders := make([]string, len(ids))
	args := make([]interface{}, len(ids))
	for i, id := range ids {
		placeholders[i] = "?"
		args[i] = id
	}

	query := "DELETE FROM tasks WHERE id IN (" + makePlaceholders(len(ids)) + ")"
	result, err := db.Exec(query, args...)
	if err != nil {
		return 0, err
	}

	deletedCount, err := result.RowsAffected()
	return deletedCount, err
}
// GetWorkerByNameAndIP 根据名称和IP地址获取Worker
func (db *DB) GetWorkerByNameAndIP(name, ip string) (*models.Worker, error) {
	var w models.Worker
	var systemInfoStr, diskInfoStr, fioVersionStr, fioStatusStr sql.NullString
	var lastHeartbeat, createdAt sql.NullTime
	err := db.QueryRow("SELECT id, name, ip_address, status, heartbeat_status, last_heartbeat, created_at, consecutive_heartbeat_failures, consecutive_heartbeat_success, system_info, disk_info, fio_version, fio_status FROM workers WHERE name = ? AND ip_address = ?", name, ip).
		Scan(&w.ID, &w.Name, &w.IPAddress, &w.Status, &w.HeartbeatStatus, &lastHeartbeat, &createdAt, &w.ConsecutiveHeartbeatFailures, &w.ConsecutiveHeartbeatSuccess, &systemInfoStr, &diskInfoStr, &fioVersionStr, &fioStatusStr)
	if err != nil {
		return nil, err
	}
	
	// 处理时间字段
	if lastHeartbeat.Valid {
		w.LastHeartbeat = lastHeartbeat.Time
	}
	if createdAt.Valid {
		w.CreatedAt = createdAt.Time
	}
	
	// 解析系统信息
	if systemInfoStr.Valid && systemInfoStr.String != "" {
		var systemInfo models.SystemInfo
		if err := json.Unmarshal([]byte(systemInfoStr.String), &systemInfo); err == nil {
			w.SystemInfo = &systemInfo
		}
	}
	
	// 解析磁盘信息
	if diskInfoStr.Valid && diskInfoStr.String != "" {
		var diskInfo []models.DiskInfo
		if err := json.Unmarshal([]byte(diskInfoStr.String), &diskInfo); err == nil {
			w.DiskInfo = diskInfo
		}
	}
	
	// 设置 fio 版本信息
	if fioVersionStr.Valid {
		w.FIOVersion = fioVersionStr.String
	}
	
	// 设置 fio 状态信息
	if fioStatusStr.Valid {
		w.FIOStatus = fioStatusStr.String
	}
	
	return &w, nil
}

// CheckWorkerHeartbeatTimeouts 检查所有Worker的心跳超时情况
func (db *DB) CheckWorkerHeartbeatTimeouts() error {
	// 开始一个事务来确保一致性
	tx, err := db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback() // 如果正常完成则提交，否则回滚

	rows, err := tx.Query("SELECT id, last_heartbeat, consecutive_heartbeat_failures FROM workers")
	if err != nil {
		return err
	}
	defer rows.Close()

	// 收集需要更新的Worker列表
	var workersToUpdate []struct {
		id           string
		status       string
		heartbeatStatus string
		failures     int
	}
	
	for rows.Next() {
		var id string
		var lastHeartbeat time.Time
		var consecutiveFailures int
		err := rows.Scan(&id, &lastHeartbeat, &consecutiveFailures)
		if err != nil {
			continue
		}

		// 如果最后心跳时间超过90秒（3次心跳间隔30秒），标记为心跳失败
		if time.Since(lastHeartbeat) > 90*time.Second {
			// 增加失败计数，重置成功计数
			newFailures := consecutiveFailures + 1
			var newHeartbeatStatus string
			
			// 如果失败次数达到3次，更改心跳状态为error
			if newFailures >= 3 {
				newHeartbeatStatus = "error"
				newFailures = 0 // 重置失败计数
			} else if newFailures >= 1 {
				// 如果失败次数至少1次但少于3次，则标记为警告
				newHeartbeatStatus = "warning"
			} else {
				newHeartbeatStatus = "healthy"
			}
			
			workersToUpdate = append(workersToUpdate, struct {
				id           string
				status       string
				heartbeatStatus string
				failures     int
			}{
				id: id,
				status: "offline",
				heartbeatStatus: newHeartbeatStatus,
				failures: newFailures,
			})
		} else {
			// 如果心跳在正常范围内，重置失败计数
			workersToUpdate = append(workersToUpdate, struct {
				id           string
				status       string
				heartbeatStatus string
				failures     int
			}{
				id: id,
				status: "", // 空字符串表示不需要更新status
				heartbeatStatus: "", // 空字符串表示不需要更新heartbeat_status
				failures: 0, // 重置失败计数
			})
		}
	}

	// 提交原始事务
	if err := tx.Commit(); err != nil {
		return err
	}

	// 更新检测到需要更改的Worker
	for _, w := range workersToUpdate {
		if w.status != "" && w.heartbeatStatus != "" {
			// 需要更新状态和心跳状态（超时情况）
			statement := `
			UPDATE workers 
			SET status = ?, heartbeat_status = ?, consecutive_heartbeat_failures = ?
			WHERE id = ?
			`
			if _, err := db.Exec(statement, w.status, w.heartbeatStatus, w.failures, w.id); err != nil {
				// 记录错误但继续处理其他Worker
				continue
			}
		} else if w.status == "" && w.heartbeatStatus == "" && w.failures == 0 {
			// 只需要重置失败计数（正常情况）
			statement := `
			UPDATE workers 
			SET consecutive_heartbeat_failures = 0
			WHERE id = ?
			`
			if _, err := db.Exec(statement, w.id); err != nil {
				// 记录错误但继续处理其他Worker
				continue
			}
		} else if w.failures == 0 {
			// 更新失败计数为0，但不改变状态
			statement := `
			UPDATE workers 
			SET consecutive_heartbeat_failures = 0
			WHERE id = ?
			`
			if _, err := db.Exec(statement, w.id); err != nil {
				// 记录错误但继续处理其他Worker
				continue
			}
		}
	}

	return nil
}